import argparse
import json
import os
import pandas as pd
from utils import read_fasta

def report_dataset(args):
    # fill information
    proteins = os.listdir(args.dataset_input)
    for protein in proteins:
        cur_dir = os.path.join(args.dataset_input, protein)
        name = protein
        sequence = read_fasta(os.path.join(cur_dir, f"{name}.fasta"))
        mutation_info = {"name": name, "seq": sequence}
        curr_seq_len = len(sequence)
        mutation_info["lenth"] = curr_seq_len
        # save info to json
        info_path = os.path.join(cur_dir, f"{name}.info.json")
        with open(info_path, "w", encoding='utf-8') as f:
            f.write(json.dumps(mutation_info))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset_input", type=str, default="data/evaluation")
    
    args = parser.parse_args()
    
    report_dataset(args)